"""
数据预处理脚本 - 准备继续训练数据

将 output_with_answers.csv 转换为训练格式，并与原有数据合并
"""
import pandas as pd
from sklearn.model_selection import train_test_split
import os

def prepare_continue_training_data():
    """准备继续训练的数据集"""

    print("=" * 70)
    print("📊 准备继续训练数据")
    print("=" * 70)

    # 1. 读取新增的非拒答数据
    print("\n1. 读取新增非拒答数据...")
    new_data_path = "../data/output_with_answers.csv"
    new_df = pd.read_csv(new_data_path)
    print(f"   新增数据: {len(new_df)} 条")

    # 2. 转换为训练格式 (text, label)
    # 非拒答内容标记为 label=0 (安全)
    print("\n2. 转换数据格式...")
    new_train_df = pd.DataFrame({
        'text': new_df['回答'],
        'label': 0  # 非拒答内容标记为安全
    })

    # 去除空值和重复项
    new_train_df = new_train_df.dropna()
    new_train_df = new_train_df.drop_duplicates()
    print(f"   清洗后数据: {len(new_train_df)} 条")

    # 3. 读取原有训练数据
    print("\n3. 读取原有训练数据...")
    old_train_path = "../data/total_train.csv"
    old_train_df = pd.read_csv(old_train_path)
    print(f"   原有训练集: {len(old_train_df)} 条")
    print(f"   - label=0 (安全): {len(old_train_df[old_train_df['label']==0])} 条")
    print(f"   - label=1 (危险): {len(old_train_df[old_train_df['label']==1])} 条")

    # 4. 合并数据
    print("\n4. 合并数据...")
    combined_df = pd.concat([old_train_df, new_train_df], ignore_index=True)
    combined_df = combined_df.drop_duplicates()
    print(f"   合并后总数: {len(combined_df)} 条")
    print(f"   - label=0 (安全): {len(combined_df[combined_df['label']==0])} 条")
    print(f"   - label=1 (危险): {len(combined_df[combined_df['label']==1])} 条")

    # 5. 重新划分训练集和验证集 (保持80/20比例)
    print("\n5. 重新划分数据集...")
    train_df, val_df = train_test_split(
        combined_df,
        test_size=0.15,  # 85% 训练, 15% 验证
        random_state=42,
        stratify=combined_df['label']  # 保持类别比例
    )

    print(f"   新训练集: {len(train_df)} 条")
    print(f"   新验证集: {len(val_df)} 条")

    # 6. 保存数据
    print("\n6. 保存数据集...")
    output_dir = "../data"

    # 保存为新文件名，避免覆盖原始数据
    train_output = os.path.join(output_dir, "continued_train.csv")
    val_output = os.path.join(output_dir, "continued_val.csv")

    train_df.to_csv(train_output, index=False, encoding='utf-8')
    val_df.to_csv(val_output, index=False, encoding='utf-8')

    print(f"   ✅ 训练集已保存: {train_output}")
    print(f"   ✅ 验证集已保存: {val_output}")

    # 7. 统计信息
    print("\n" + "=" * 70)
    print("📈 数据统计摘要")
    print("=" * 70)
    print(f"原有训练数据: {len(old_train_df)} 条")
    print(f"新增非拒答数据: {len(new_train_df)} 条")
    print(f"合并后总数: {len(combined_df)} 条")
    print(f"\n新训练集: {len(train_df)} 条")
    print(f"  - label=0: {len(train_df[train_df['label']==0])} ({len(train_df[train_df['label']==0])/len(train_df)*100:.1f}%)")
    print(f"  - label=1: {len(train_df[train_df['label']==1])} ({len(train_df[train_df['label']==1])/len(train_df)*100:.1f}%)")
    print(f"\n新验证集: {len(val_df)} 条")
    print(f"  - label=0: {len(val_df[val_df['label']==0])} ({len(val_df[val_df['label']==0])/len(val_df)*100:.1f}%)")
    print(f"  - label=1: {len(val_df[val_df['label']==1])} ({len(val_df[val_df['label']==1])/len(val_df)*100:.1f}%)")
    print("\n✅ 数据准备完成！")

if __name__ == "__main__":
    prepare_continue_training_data()
